# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import json
from typing import List, Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.okex_futures.kr_rest.futures_constants import contract_type_list
from coin.exchange.okex.parse_util import parse_okex_msg2


class OkexFuturesSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None,
               subscribe_ticker=True):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)
    self._register_url('wss://real.okex.com:10440/websocket?compress=true')

    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

    self._depth_msg_count = {symbol: 0 for symbol in self._symbol}
    self._last_pong_time = None
    self._subscribe_ticker = subscribe_ticker

  @gen.coroutine
  def on_open(self):
    self._logger.debug('Connected')
    self._ioloop.add_callback(self._send_subscribe_request)

  def _decode(self, ws_msg):
    return parse_okex_msg2(ws_msg)

  @gen.coroutine
  def _send_subscribe_request(self):
    underlying_set = set()
    for symbol in self._symbol:
      assert symbol[3] == '_'
      underlying = symbol[0:3]
      expiration = symbol[4:]
      assert expiration in contract_type_list

      self._logger.debug('Subscribing %s' % symbol)

      if underlying not in underlying_set:
        underlying_set.add(underlying)
        yield self._ws.write_message(
            json.dumps({
                'event': 'addChannel', 'channel': 'ok_sub_futureusd_%s_index' % underlying
            }))

      yield self._ws.write_message(
          json.dumps({
              'event': 'addChannel',
              'channel': 'ok_sub_futureusd_%s_depth_%s' % (underlying, expiration)
          }))
      yield self._ws.write_message(
          json.dumps({
              'event': 'addChannel',
              'channel': 'ok_sub_futureusd_%s_depth_%s_20' % (underlying, expiration)
          }))
      yield self._ws.write_message(
          json.dumps({
              'event': 'addChannel',
              'channel': 'ok_sub_futureusd_%s_trade_%s' % (underlying, expiration)
          }))
      if self._subscribe_ticker:
        yield self._ws.write_message(
            json.dumps({
                'event': 'addChannel',
                'channel': 'ok_sub_futureusd_%s_ticker_%s' % (underlying, expiration)
            }))
      yield gen.sleep(0.5)  # Sleep 0.5 second before next request

    self._publish_topic_event('START')
    self._logger.info('Subscribe done')

    self._last_pong_time = datetime.datetime.now()
    while self._state == self.State.RUNNING:
      if (datetime.datetime.now() - self._last_pong_time >= datetime.timedelta(minutes=2)):
        self._logger.error('No pong has been received since %s. Closing...'
                           % str(self._last_pong_time))
        self._notify_close()
        break

      self._logger.debug('ping')
      yield self._ws.write_message('{"event":"ping"}')
      yield gen.sleep(30)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_futureusd_depth_fullbook = topic_map.search_topic_id(
        'okex_ws_futureusd_depth_fullbook_%s%s' % (group_id, worker_id))[0]
    topic_futureusd_depth_diff = topic_map.search_topic_id('okex_ws_futureusd_depth_diff_%s%s' %
                                                           (group_id, worker_id))[0]
    topic_futureusd_depth20 = topic_map.search_topic_id('okex_ws_futureusd_depth_20_%s%s' %
                                                        (group_id, worker_id))[0]
    topic_futureusd_trade = topic_map.search_topic_id('okex_ws_futureusd_trade_%s%s' %
                                                      (group_id, worker_id))[0]
    topic_futureusd_index = topic_map.search_topic_id('okex_ws_futureusd_index_%s%s' %
                                                      (group_id, worker_id))[0]
    topic_futureusd_ticker = topic_map.search_topic_id('okex_ws_futureusd_ticker_%s%s' %
                                                       (group_id, worker_id))[0]

    self._publisher = {
        'depth_fullbook': queue.get_writer(topic_id=topic_futureusd_depth_fullbook[1]),
        'depth_diff': queue.get_writer(topic_id=topic_futureusd_depth_diff[1]),
        'depth_20': queue.get_writer(topic_id=topic_futureusd_depth20[1]),
        'trade': queue.get_writer(topic_id=topic_futureusd_trade[1]),
        'index': queue.get_writer(topic_id=topic_futureusd_index[1]),
        'ticker': queue.get_writer(topic_id=topic_futureusd_ticker[1])
    }

    self._enable_topic_event_publisher(queue,
                                       'okex_ws_topic_event',
                                       [
                                           topic_futureusd_depth_fullbook[0],
                                           topic_futureusd_depth_diff[0],
                                           topic_futureusd_depth20[0],
                                           topic_futureusd_trade[0],
                                           topic_futureusd_index[0],
                                           topic_futureusd_ticker[0]
                                       ])
    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def _on_add_channel(self, timestamp: int, message: dict):
    try:
      data = message['data']
      channel = data['channel']
      if data['result']:
        self._logger.debug('Channel subscribed: %s' % channel)
      else:
        self._logger.error('Failed to subscriber channel: %s' % channel)
    except KeyError:
      self._logger.error('Invalid message\n%s' % message)

  def _on_depth(self, timestamp: int, channel: str, message: dict):
    currency = channel[17:20]
    if channel.endswith('this_week'):
      symbol = '%s_this_week' % currency
    elif channel.endswith('next_week'):
      symbol = '%s_next_week' % currency
    elif channel.endswith('quarter'):
      symbol = '%s_quarter' % currency
    else:
      self._logger.error('Invalid channel: %s' % channel)
      return

    count = self._depth_msg_count[symbol]
    self._depth_msg_count[symbol] += 1
    if count == 0:
      self._publisher['depth_fullbook'].write(timestamp, message)
    else:
      self._publisher['depth_diff'].write(timestamp, message)

  def on_message(self, timestamp: int, message: Union[dict, list]):
    messages = to_list(message)
    for msg in messages:
      if self._state != self.State.RUNNING:
        break

      if msg is None:
        continue

      channel = msg.get('channel', '')
      if channel.find('_depth_') != -1:
        if channel.endswith('_20'):
          self._publisher['depth_20'].write(timestamp, msg)
        else:
          self._on_depth(timestamp, channel, msg)

      elif channel.find('_trade_') != -1:
        # TODO: okex sends old trades at the beginning. Ignore first
        # messages.
        if self._last_pong_time is not None:
          self._publisher['trade'].write(timestamp, msg)

      elif channel.endswith('_index'):
        self._publisher['index'].write(timestamp, msg)

      elif channel.find('_ticker_') != -1:
        self._publisher['ticker'].write(timestamp, msg)

      elif channel == 'addChannel':
        self._on_add_channel(timestamp, msg)

      elif message.get('event', None) == 'pong':
        self._last_pong_time = datetime.datetime.now()
        self._logger.debug('pong')

      else:
        self._logger.error('Unknown message: %s' % str(msg))
